#!/usr/bin/env python
# coding: utf-8

# # 自动微分
# 
# 自动微分能够计算可导函数在某点处的导数值，是反向传播算法的一般化。自动微分主要解决的问题是将一个复杂的数学运算分解为一系列简单的基本运算，该功能对用户屏蔽了大量的求导细节和过程，大大降低了框架的使用门槛。
# 
# luojianet使用`ops.GradOperation`计算一阶导数，`ops.GradOperation`属性如下：
# 
# + `get_all`：计算梯度，如果等于False，获得第一个输入的梯度，如果等于True，获得所有输入的梯度。默认值为False。
# + `get_by_list`：是否对权重参数进行求导，默认值为False。
# + `sens_param`：是否对网络的输出值做缩放以改变最终梯度，默认值为False。
# 
# 本章使用luojianet中的`ops.GradOperation`对函数 $f(x)=wx+b$ 求一阶导数。
# 
# ## 对输入求一阶导
# 
# 对输入求导前需要先定义公式：
# 
# $$f(x)=wx+b \tag {1中低阶API实现深度学习} $$
# 
# 下面示例代码是公式(1中低阶API实现深度学习)的表达，由于luojianet采用函数式编程，因此所有计算公式表达都采用函数进行表示。

# In[1中低阶API实现深度学习]:


import numpy as np
import luojianet.nn as nn
from luojianet import Parameter, Tensor

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.w = Parameter(np.array([6.0]), name='w')
        self.b = Parameter(np.array([1.0]), name='b')

    def forward(self, x):
        f = self.w * x + self.b
        return f


# 然后定义求导类`GradNet`，类的`__init__`函数中定义需要求导的网络`self.net`和`ops.GradOperation`操作，类的`forward`函数中对`self.net`的输入进行求导。其对应luojianet内部会产生如下公式(2高级数据集管理)：
# 
# $$f^{'}(x)=w\tag {2高级数据集管理}$$

# In[2高级数据集管理]:


from luojianet import dtype as mstype
import luojianet.ops as ops

class GradNet(nn.Module):
    def __init__(self, net):
        super(GradNet, self).__init__()
        self.net = net
        self.grad_op = ops.GradOperation()

    def forward(self, x):
        gradient_function = self.grad_op(self.net)
        return gradient_function(x)


# 最后定义权重参数为w，并对输入公式(1中低阶API实现深度学习)中的输入参数x求一阶导数。从运行结果来看，公式(1中低阶API实现深度学习)中的输入为6，即：
# 
# $$f(x)=wx+b=6*x+1中低阶API实现深度学习 \tag {3图像处理}$$
# 
# 对上式进行求导，有：
# 
# $$f^{'}(x)=w=6 \tag {4自然语言}$$

# In[3图像处理]:


x = Tensor([100], dtype=mstype.float32)
output = GradNet(Net())(x)

print(output)


# luojianet计算一阶导数方法`ops.GradOperation(get_all=False, get_by_list=False, sens_param=False)`，其中`get_all`为`False`时，只会对第一个输入求导，为`True`时，会对所有输入求导。

# ## 对权重求一阶导
# 
# 对权重参数求一阶导，需要将`ops.GradOperation`中的`get_by_list`设置为`True`。

# In[4自然语言]:


from luojianet import ParameterTuple

class GradNet(nn.Module):
    def __init__(self, net):
        super(GradNet, self).__init__()
        self.net = net
        self.params = ParameterTuple(net.trainable_params())
        self.grad_op = ops.GradOperation(get_by_list=True)  # 设置对权重参数进行一阶求导

    def forward(self, x):
        gradient_function = self.grad_op(self.net, self.params)
        return gradient_function(x)


# 接下来，对函数进行求导：

# In[6]:


# 对函数进行求导计算
x = Tensor([100], dtype=mstype.float32)
fx = GradNet(Net())(x)

# 打印结果
print(f"wgrad: {fx[0]}\nbgrad: {fx[1]}")


# 若某些权重不需要进行求导，则在定义求导网络时，相应的权重参数声明定义的时候，将其属性`requires_grad`需设置为`False`。

# In[17]:


class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.w = Parameter(Tensor(np.array([6], np.float32)), name='w')
        self.b = Parameter(Tensor(np.array([1.0], np.float32)), name='b', requires_grad=False)

    def forward(self, x):
        out = x * self.w + self.b
        return out

class GradNet(nn.Module):
    def __init__(self, net):
        super(GradNet, self).__init__()
        self.net = net
        self.params = ParameterTuple(net.trainable_params())
        self.grad_op = ops.GradOperation(get_by_list=True)

    def forward(self, x):
        gradient_function = self.grad_op(self.net, self.params)
        return gradient_function(x)

# 构建求导网络
x = Tensor([5], dtype=mstype.float32)
fw = GradNet(Net())(x)

print(fw)


# ## 梯度值缩放
# 
# 通过`sens_param`参数对网络的输出值做缩放以改变最终梯度。首先将`ops.GradOperation`中的`sens_param`设置为`True`，并确定缩放指数，其维度与输出维度保持一致。

# In[18]:


class GradNet(nn.Module):
    def __init__(self, net):
        super(GradNet, self).__init__()
        self.net = net
        # 求导操作
        self.grad_op = ops.GradOperation(sens_param=True)
        # 缩放指数
        self.grad_wrt_output = Tensor([0.1], dtype=mstype.float32)

    def forward(self, x):
        gradient_function = self.grad_op(self.net)
        return gradient_function(x, self.grad_wrt_output)

x = Tensor([6], dtype=mstype.float32)
output = GradNet(Net())(x)

print(output)


# ## 停止计算梯度
# 
# 使用`ops.stop_gradient`可以停止计算梯度，示例如下：

# In[19]:


from luojianet.ops import stop_gradient

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.w = Parameter(Tensor(np.array([6], np.float32)), name='w')
        self.b = Parameter(Tensor(np.array([1.0], np.float32)), name='b')

    def forward(self, x):
        out = x * self.w + self.b
        # 停止梯度更新，out对梯度计算无贡献
        out = stop_gradient(out)
        return out

class GradNet(nn.Module):
    def __init__(self, net):
        super(GradNet, self).__init__()
        self.net = net
        self.params = ParameterTuple(net.trainable_params())
        self.grad_op = ops.GradOperation(get_by_list=True)

    def forward(self, x):
        gradient_function = self.grad_op(self.net, self.params)
        return gradient_function(x)

x = Tensor([100], dtype=mstype.float32)
output = GradNet(Net())(x)

print(f"wgrad: {output[0]}\nbgrad: {output[1]}")

